# -*- coding: utf-8 -*-
# -*- coding: utf-8 -*-

from pymongo import UpdateOne, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes, get_trading_dates
from util.database import DB_CONN
import pandas as pd
import numpy as np

"""
实现市盈率因子的计算和保存
"""


class SP_TTM_Factor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='pe')

    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值，并保存到数据库中
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        dm = DataModule()
        dates = get_trading_dates(begin_date, end_date)
        codes = ['601808']
        update_requests = []
        for code in codes:
            data = DB_CONN['lrb'].find({'code': code})
            cc = [asa['date'] for asa in data]
            for date in dates:
                for asa in cc:
                    if date < asa:
                        continue
                    else:
                        reportdat = asa
                        break

                df = DB_CONN['mainreport'].find_one({'code': code, 'date': reportdat})

                NPG = df['gsjlrtbzz']

                basic = DB_CONN['basic'].find_one({'code': code, 'date': date})
                close = dm.get_k_data(code=code, begin_date=date, end_date=date)

                TTM = int(basic['totals'] * close['close'])
                NPG_TTM = (NPG / TTM)

                update_requests.append(
                    UpdateOne(
                        {'code': code, 'date': date},
                        {'$set': {'code': code, 'date': date, 'NPG_TTM': NPG_TTM}}, upsert=True))
                print(update_requests)
                break
            break


SP_TTM_Factor().compute('2018-01-02', '2018-09-05')
